package cn.lgwen.spark.ml.learning.kaggle;

import org.apache.spark.api.java.function.MapFunction;
import org.apache.spark.ml.classification.RandomForestClassificationModel;
import org.apache.spark.ml.classification.RandomForestClassifier;
import org.apache.spark.ml.evaluation.MulticlassClassificationEvaluator;
import org.apache.spark.ml.feature.Normalizer;
import org.apache.spark.ml.feature.VectorAssembler;
import org.apache.spark.ml.linalg.VectorUDT;
import org.apache.spark.sql.Dataset;
import org.apache.spark.sql.Row;
import org.apache.spark.sql.RowFactory;
import org.apache.spark.sql.SparkSession;
import org.apache.spark.sql.catalyst.encoders.RowEncoder;
import org.apache.spark.sql.types.DataTypes;
import org.apache.spark.sql.types.Metadata;
import org.apache.spark.sql.types.StructField;
import org.apache.spark.sql.types.StructType;

/**
 * 2021/4/29
 *
 * @author aven@didiglobal.com
 */
public class RedWineQuality {

    public static void main(String[] args) {
        SparkSession spark = SparkSession
                .builder().master("local[*]")
                .appName("RedWineQuality")
                .getOrCreate();

        Dataset<Row> dataset = trainDataset(spark);
        dataset.registerTempTable("red_wine");

        Normalizer normalizer = new Normalizer()
                .setInputCol("features")
                .setOutputCol("normFeatures")
                .setP(1.0);
        dataset = normalizer.transform(dataset);
        dataset.show(10);

        Dataset<Row>[] splits = dataset.randomSplit(new double[]{0.5, 0.5});

        RandomForestClassifier rf = new RandomForestClassifier()
                .setLabelCol("label")
                .setFeaturesCol("normFeatures");
        RandomForestClassificationModel model = rf.fit(splits[0]);


        Dataset<Row> testData = model.transform(splits[1]);
        MulticlassClassificationEvaluator evaluator = new MulticlassClassificationEvaluator()
                .setPredictionCol("prediction")
                .setLabelCol("label");



        double evaluate = evaluator.evaluate(testData);

        System.out.println(evaluate);


    }



    public static Dataset<Row> trainDataset(SparkSession spark) {
        Dataset<Row> trainData = spark.read().format("csv").option("header", true).option("inferSchema", true)
                .load("/Users/didi/Develop/data/winequality-red.csv");

        VectorAssembler vectorAssem4 = new VectorAssembler()
                .setInputCols(new String[]{"fixed acidity", "volatile acidity", "citric acid", "residual sugar",
                        "chlorides", "free sulfur dioxide", "total sulfur dioxide", "density",
                        "pH","sulphates","alcohol"
                }).
                        setOutputCol("features").setHandleInvalid("keep");

        StructType AgeComplementSchema = new StructType(new StructField[]{
                new StructField("label", DataTypes.DoubleType, true, Metadata.empty()),
                new StructField("features", new VectorUDT(), false, Metadata.empty()),
        });

        //return vectorAssem4.transform(trainData);
        return vectorAssem4.transform(trainData).map(new MapFunction<Row, Row>() {
            @Override
            public Row call(Row row) throws Exception {
                return RowFactory.create((double)row.getInt(11), row.get(12));
            }
        }, RowEncoder.apply(AgeComplementSchema));


    }
}
